Robust Spectrotemporal Decomposition by Iteratively Reweighted Least Squares

Python notebook written by: Armen Gharibans

Reference: Ba, D., Babadi, B., Purdon, P. L., & Brown, E. N. (2014). Robust spectrotemporal decomposition by iteratively reweighted least squares. Proceedings of the National Academy of Sciences, 111(50), E5336-E5345.

Toy Example

This example points out the limitations of classical techniques in analyzing time series data. We simulated noisy observations from the linear combination of two amplitude-modulated signals using the following equation:

$y_{t}=10 \cos^{8}\left(2\pi f_{o}t\right) \sin\left(2\pi f_{1}t\right)+10\exp\left(4\frac{t-T}{T}\right)\cos\left(2\pi f_{2}t\right)+v_{t}, \quad$ for $\enspace 0\leq t \leq T$

where $f_{0}=0.04$ Hz, $f_{1}=10$ Hz, $f_{2}=11$ Hz, $T=600$ s, and $\left(v_{t}\right)_{t=1}^{T}$ is independent, identically distributed, zero-mean Gaussian noise with variance set to acheive a signal-to-noise ratio (SNR) of 5 dB.

The simulated data consists of a 10 Hz oscillation whos amplitude is modulated by slow 0.04 Hz oscillation, and an exponentially gorwing 11 Hz oscillation.

The following is a spectrogram of the simulated signal that highlights the limitations of classical frequency analysis. The analysis is unable to resolve the closely spaced signals of 10 and 11 Hz in the frequency domain.

Robust Spectral Decomposition

The State-Space Model

In this analysis, we will consider a signal $y_{t}$ that is obtained by sampling a noisy, continuous-time signal at rate $f_{s}$ (above the Nyquist rate).

• $y_{t}\rightarrow$ discrete-time signal
• $t = 1,2,...,T\rightarrow$ samples
• $f_{s}\rightarrow$ sampling rate
• $W\rightarrow$ arbitrary window length
• $N \triangleq \frac{T}{W}\rightarrow$ number of windows
• $y_{n} \triangleq \left(y_{\left(n-1\right)W+1},y_{\left(n-1\right)W+2},...,y_{nW}\right)'$ for $n=1,2,...,N$

Consider the following spectrotemporal representation of $y_{n}$ as:

$$y_{n} = \tilde{F}_{n}\tilde{x}_{n} + v_{n}$$

where,

• $\left ( \tilde{F}_{n} \right )_{l,k} \triangleq \exp \left( j2\pi \left( \left(n-1 \right) W+l \right) \frac{\left( k-1 \right)}{K} \right)$ for $l=1,2,...,W$ and $k=1,2,...K$
• $\tilde{x}_{n} \triangleq \left(\tilde{x}_{n,1},\tilde{x}_{n,2},...,\tilde{x}_{n,K} \right)' $
• $v_{n}\rightarrow$ independent, identically distributed, additive zero-mean Gaussian noise

Equivalently, we can define the linear observation model over a real vector space as follows:

$$y_{n} = F_{n}x_{n}+v_{n}$$

where,

• $\left ( F_{n} \right )_{l,k} \triangleq \cos \left( 2\pi \left( \left(n-1 \right) W+l \right) \frac{\left( k-1 \right)}{K} \right)$ for $l = 1,2,...,W$ and $k=1,2,...\frac{K}{2}$
• $\left ( F_{n} \right )_{l,k+K/2} \triangleq \sin \left( 2\pi \left( \left(n-1 \right) W+l \right) \frac{\left( k-1 \right)}{K} \right)$ for $l = 1,2,...,W$ and $k=1,2,...\frac{K}{2}$
• $x_{n} \triangleq \left(x_{n,1},x_{n,2},...,x_{n,K} \right)' $

The objective is to compute an estimate $\hat{x}$ of $x$ given the data $y$. The component-wise magnitude-squared of $\hat{x}$ gives an estimate of the magnitude spectrum of $y$. By treating $\left(x_{n}\right)_{n=1}^{N}$ as a sequence of random variables and carefully selecting a prior distribution, a stochastic continuity constraint can be established across time. By imposing a model on the components $\left(x_{n,k}\right)_{k=1}^{K}$ for each $n = 1,2,..,N$, sparsity is enforced in the frequency domain. The stochastic continuity constraint can be expressed in the form of the first-order difference equation:

$$x_{n} = x_{n-1} + w_{n}$$

where $w = \left( w_{1}',w_{2}',...,w_{N}'\right)'$ is a random vector. The following joint prior probability density function is used to enforce sparsity in the frequency domain and smoothness in time: $$\log p_{1}\left(w_{1},w_{2},...,w_{N}\right) = -\alpha \sum_{k=1}^{K} \left( \sum_{n=1}^{N} w_{n,k}^{2} + \epsilon^{2} \right)^\frac{1}{2}+c_{1}$$ where $\alpha > 0$ is a constant and $\epsilon > 0$ is a small constant.

The Inverse Solution

Bayesian estimation is used to compute the robust spectral decomposition of $y$, where the posterior density of $x$ given $y$ fully characterizes the space of inverse solutions. This is computed by solving the following MAP estimation problem:

$$\max_{x_{1},...x_{N}} -\sum_{n=1}^{N} \frac{1}{2\sigma^{2}} \left \| y_{n} - F_{n}x_{n} \right \|^{2}_{2} + f\left(x_{1},x_{2},...,x_{N}\right)$$

where $f\left(x_{1},x_{2},...,x_{N}\right) \triangleq \log p_{i} \left(x_{1}-x_{0},x_{2}-x_{1},...,x_{N}-x_{N-1}\right)$. This is a strictly concave optimization problem that can be solved using standard techniques. However, these techniques do not scale well with $N$ becuase of the batch nature of the problem.

Kalman Filter

The Kalman filter solves the least-squares estimation problem recursively, and in a computationally efficient manner. The algorithm works in a two-step process. In the prediciton step, the Kalman filter produces estimates of the current state variables, along with their uncertainties. Once the outcome of the next measurement (corrupted with some amount of error, including random noise) is observed, these estimates are updated with more weight being given to estimates with higher certainty.

Algorithm:

Initial Conditions:

• $x_{0\mid 0} = \left(0,...,0\right)' \in \mathbb{R}^{K}$
• $\Sigma_{0\mid 0} = I_{k} \in \mathbb{R}^{KK}$

Filter at time $n=1,2,...,N$:

• $x_{n\mid n-1}=x_{n-1 \mid n-1}$
• $\Sigma_{n\mid n-1}=\Sigma_{n-1\mid n-1}+Q^{\left(l\right)}$
• $K_{n}=\Sigma_{n\mid n-1}F_{n}^{H}\left(F_{n}\Sigma_{n\mid n-1}F_{n}^{H}+\sigma^{2}I\right)^{-1}$
• $x_{n\mid n}=x_{n\mid n-1}+K_{n}\left(y_{n}-F_{n}x_{n\mid n-1}\right)$
• $\Sigma_{n\mid n}=\Sigma_{n\mid n-1}-K_{n}F_{n}\Sigma_{n\mid n-1}$

Kalman Smoother

The Kalman filter is designed for real-time applications. It estimates the properties of a system at a given time using measurments of the system up to that time. However, when a real-time estimate is not needed, the Kalman filter effectively throws away half of the measurement data. The Kalman smoother is an extension of the Kalman filter that uses measurement information from after the time at which state estimates are required as well as before that time.

Algorithm

Smoother at time $n=N-1,N-2,...,1:$

• $B_{n}=\Sigma_{n\mid n}\Sigma_{n+1\mid n}^{-1}$
• $x_{n\mid N}=x_{n\mid n}+B_{n}\left(x_{n+1\mid N}-x_{n+1\mid n}\right)$
• $\Sigma_{n\mid N}=\Sigma_{n\mid n}+B_{n}\left(\Sigma_{n+1\mid N}-\Sigma_{n+1\mid n}\right)B_{n}^{H}$

IRLS Algorithm for Spectrotemporal Pursuit

The solution to the optimization problem can be obtained as the limit of a sequence $\left(\hat{x}^\left(l\right)\right)_{l=0}^{\infty}$ whose $l^{th}$ element is the solution to the Guassian MAP estimation problem (constrained least-squares program) of the form:

$$ \max_{x_{1},...,x{N}} -\sum_{n=1}^{N} \frac{1}{2\sigma^{2}} \left \| y_{n} - F_{n}x_{n} \right \|^{2}_{2} - \sum_{k=1}^{K} \sum_{n=1}^{N} \frac{\left(x_{n,k}-x_{n-1,k}\right)^{2}}{2\left(Q^{\left(l\right)}\right)_{k,k}} $$

where, $$ \left(Q^{\left(l\right)}\right)_{k,k} = \frac{\left(\sum_{n=1}^{N}\left(\hat{x}_{n,k}^{\left(l-1\right)}-\hat{x}_{n-1,k}^{\left(l-1\right)} \right)+\epsilon^{2}\right)^{\frac{1}{2}}}{\alpha} \rightarrow K\times K \textrm{ diagonal matrix}$$

This is a quadratic program with strictly concave objective function and block-tridiagonal Hessian. It can be solved iteratively using the following steps:

Input:

• $y \rightarrow$ observations
• $\hat{x}^{\left(0\right)} \in \mathbb{R}^{KN}\rightarrow$ initial guess
• $Q^{\left(0\right)} \rightarrow$ initial state-noise covariance
• $x_{0\mid0}, \Sigma_{0\mid0} \rightarrow$ intial conditions
• $tol \in \left(0,0.01\right) \rightarrow$ tolerance
• $L_{max} \in \mathbb{N}^{+} \rightarrow$ maximum number of iterations

Step 0. Initialize iteration number $l$ to 1

Step 1. Kalman Filter at time $n=1,2,...,N$:

• $x_{n\mid n-1}=x_{n-1 \mid n-1}$
• $\Sigma_{n\mid n-1}=\Sigma_{n-1\mid n-1}+Q^{\left(l\right)}$
• $K_{n}=\Sigma_{n\mid n-1}F_{n}^{H}\left(F_{n}\Sigma_{n\mid n-1}F_{n}^{H}+\sigma^{2}I\right)^{-1}$
• $x_{n\mid n}=x_{n\mid n-1}+K_{n}\left(y_{n}-F_{n}x_{n\mid n-1}\right)$
• $\Sigma_{n\mid n}=\Sigma_{n\mid n-1}-K_{n}F_{n}\Sigma_{n\mid n-1}$

Step 2. Smoother at time $n=N-1,N-2,...,1$:

• $B_{n}=\Sigma_{n\mid n}\Sigma_{n+1\mid n}^{-1}$
• $x_{n\mid N}=x_{n\mid n}+B_{n}\left(x_{n+1\mid N}-x_{n+1\mid n}\right)$
• $\Sigma_{n\mid N}=\Sigma_{n\mid n}+B_{n}\left(\Sigma_{n+1\mid N}-\Sigma_{n+1\mid n}\right)B_{n}^{H}$

Step3. Let $\hat{x}_{n}^{\left(l\right)}=x_{n\mid N},n=1,...,N$ and $\hat{x}^{\left(l\right)}=\left(\hat{x}_{1}^{\left(l\right)'},...,\hat{x}_{N}^{\left(l\right)'}\right)'$

Step4. Stop if $\frac{\left \| \hat{x}^{\left(l \right )}-\hat{x}^{\left(l-1 \right) } \right \|_{2}}{\left \| \hat{x}^{\left(l-1 \right) }\right \|_{2}}<tol$ or $l=L_{max}$

Step5. Let $l=l+1$, and update the state covariance $Q_n^{\left(l\right)}$

Step6. Go back to Step 1

Output: $\hat{x}^{\left(L\right)}$, where $L\leq L_{max}$ is the number of the last iteration of the algorithm

Spectrotemporal pursuit gives the sparse, more compact representation that better characterizes the equation simulated in this toy example. We are able to recover the two frequencies (10 and 11 Hz) as well their temporal modulation. The estimate is also significantly denoised relative to the STFT spectrogram.